# models/dynamic_builder.py

import torch.nn as nn
from models.blocks import ResidualBlock

class DynamicModelBuilder:
    @staticmethod
    def build_layer(layer_config, input_dim):
        layer_type = layer_config["type"]
        params = layer_config.get("params", {})
        activation = layer_config.get("activation")

        if layer_type == "Linear":
            layer = nn.Linear(in_features=input_dim, out_features=params["out_features"])
            output_dim = params["out_features"]
        elif layer_type == "Dropout":
            layer = nn.Dropout(**params)
            output_dim = input_dim
        elif layer_type == "BatchNorm1d":
            layer = nn.BatchNorm1d(input_dim)
            output_dim = input_dim
        elif layer_type == "ResidualBlock":
            layer = ResidualBlock(input_dim, params["hidden_dim"])
            output_dim = input_dim
        else:
            raise ValueError(f"Unsupported layer type: {layer_type}")

        act = None
        if activation:
            act = getattr(nn, activation)()

        return layer, act, output_dim
